# LeCrunch
# Copyright (C) 2010 Anthony LaTorre 
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.

import os
import sys
import numpy
import ROOT

dir = os.path.dirname(os.path.realpath(__file__))

ROOT.gROOT.ProcessLine('.L ' + dir + '/ScopeEvent.C+')

class RootFile:
    def __init__(self, filename, mode='r'):
        self.mode = mode
        if self.mode == 'w':
            self.f = ROOT.TFile(filename, 'recreate')
            self.T = ROOT.TTree('T', 'Scope Events')
            self.ev = ROOT.ScopeEvent()
            self.T.Branch('event', 'ScopeEvent', ROOT.AddressOf(self.ev))
        if self.mode == 'r':
            self.f = ROOT.TFile(filename, 'read')
            if self.f.IsZombie():
                raise Exception("file '%s' is a zombie' % filename")
            self.T = self.f.Get('T')
            self.ev = ROOT.ScopeEvent()
            self.T.SetBranchAddress('event', self.ev)

    def __len__(self):
        return int(self.T.GetEntries())

    def getentry(self, i):
        self.T.GetEntry(i)

    def getchannels(self):
        channels = []
        if self.ev.GetC1():
            channels.append(1)
        if self.ev.GetC2():
            channels.append(2)
        if self.ev.GetC3():
            channels.append(3)
        if self.ev.GetC4():
            channels.append(4)
        return channels

    def getc1(self):
        x = numpy.ndarray((self.ev.GetC1().GetN(),), dtype=numpy.double,
                          buffer=self.ev.GetC1().GetX())
        y = numpy.ndarray((self.ev.GetC1().GetN(),), dtype=numpy.double,
                          buffer=self.ev.GetC1().GetY())
        return x, y

    def iterc1(self, n):
        for i in range(len(self)):
            self.getentry(i)
            yield self.getc1()

            if i+1 >= n:
                break

    def getc2(self):
        x = numpy.ndarray((self.ev.GetC2().GetN(),), dtype=numpy.double,
                          buffer=self.ev.GetC2().GetX())
        y = numpy.ndarray((self.ev.GetC2().GetN(),), dtype=numpy.double,
                          buffer=self.ev.GetC2().GetY())
        return x, y

    def iterc2(self, n):
        for i in range(len(self)):
            self.getentry(i)
            yield self.getc2()

            if i+1 >= n:
                break

    def getc3(self):
        x = numpy.ndarray((self.ev.GetC3().GetN(),), dtype=numpy.double,
                          buffer=self.ev.GetC3().GetX())
        y = numpy.ndarray((self.ev.GetC3().GetN(),), dtype=numpy.double,
                          buffer=self.ev.GetC3().GetY())
        return x, y

    def iterc3(self, n):
        for i in range(len(self)):
            self.getentry(i)
            yield self.getc3()

            if i+1 >= n:
                break

    def getc4(self):
        x = numpy.ndarray((self.ev.GetC4().GetN(),), dtype=numpy.double,
                          buffer=self.ev.GetC4().GetX())
        y = numpy.ndarray((self.ev.GetC4().GetN(),), dtype=numpy.double,
                          buffer=self.ev.GetC4().GetY())
        return x, y

    def iterc4(self, n):
        for i in range(len(self)):
            self.getentry(i)
            yield self.getc4()

            if i+1 >= n:
                break
        
    def setc1(self, x, y):
        self.ev.SetC1(ROOT.TGraph(len(x), x, y))

    def setc2(self, x, y):
        self.ev.SetC2(ROOT.TGraph(len(x), x, y))

    def setc3(self, x, y):
        self.ev.SetC3(ROOT.TGraph(len(x), x, y))

    def setc4(self, x, y):
        self.ev.SetC4(ROOT.TGraph(len(x), x, y))

    def fill(self):
        self.T.Fill()

    def close(self):
        if self.mode == 'w':
            self.T.GetCurrentFile().cd()
            self.T.Write()
            self.T.GetCurrentFile().Close()
